package com.example.isoftlangchainai.rag.utils;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;
import java.util.List;

/**
 * @Description: Utils工具类 初始化PgVector库
 * @Date :2025/7/16 15:11
 * @Author :chenjun
 */
public class PgVectorStoreUtils {

    /**
     * [描述]:初始化PgVector库
     */
    public static EmbeddingStore<TextSegment> initPgVector() {
            System.out.println("*********创建EmbeddingModel********");
            EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();
            System.out.println("*********通过PgVectorEmbeddingStore 连接PostgreSQL数据库（PgVector）********");
            EmbeddingStore<TextSegment> embeddingStore = PgVectorEmbeddingStore.builder()
                    .dropTableFirst(true)
                    .createTable(true)
                    .host("localhost")
                    .port(5432)
                    .user("postgres")
                    .password("postgres")
                    .database("postgres")
                    .table("favouritesport")
                    .dimension(384)
                    .build();

            System.out.println("*********TextSegment Embedding并存储EmbeddingStore********");
            // 定义文本，转换成文本段对象TextSegment
            TextSegment segment0 = TextSegment.from("我最喜欢篮球.");
            //文本转换为嵌入向量Embedding
            Embedding embedding0 = embeddingModel.embed(segment0).content();
            //嵌入向量及其对应文本存入存储PgVector库
            embeddingStore.add(embedding0, segment0);

            TextSegment segment1 = TextSegment.from("我喜欢足球.");
            Embedding embedding1 = embeddingModel.embed(segment1).content();
            embeddingStore.add(embedding1, segment1);

            TextSegment segment2 = TextSegment.from("今天天气真好.");
            Embedding embedding2 = embeddingModel.embed(segment2).content();
            embeddingStore.add(embedding2, segment2);
            System.out.println("*********完成EmbeddingStore存储********");
            return embeddingStore;
    }

    /**
     *
     * [描述]:初始化嵌入embedding向量相似性搜索
     * @param userQuery      用户问题,如：你最喜欢的运动是什么?
     * @param embeddingStore 嵌入embedding存储对象
     */
    public static void similarityQuery(String userQuery, EmbeddingStore<TextSegment> embeddingStore) {
        //创建EmbeddingModel
        EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();
        // 查询内容转换成查询向量Embedding
        Embedding queryEmbedding = embeddingModel.embed(userQuery).content();
        //构建EmbeddingSearchRequest请求，设置查询向量和最大返回结果数
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(2)
                .minScore(0.7)
                .build();
        //执行向量相似性搜索,获取匹配结果列表
        List<EmbeddingMatch<TextSegment>> relevant = embeddingStore.search(embeddingSearchRequest).matches();
        System.out.println("*********用户问题:"+  userQuery +",向量相似性搜索返回的结果集:");
        relevant.forEach(embeddingMatch -> System.out.println("embedded score:"+ embeddingMatch.score() + ", embedded Context:"+ embeddingMatch.embedded().text()));
        System.out.println("*********用户问题回复完成*********");
    }
}
